from collections import defaultdict
from math import log
import json


"""
("B", "M", "E", "S")
用于分词的几个状态：
B: Begin of token 单词(标记)的开始
M: Middle of token 单词的中间
E: End of token 单词的结尾
S: Single token 单字词
"""


class TokenizerTrainer:
    init_prob: dict = defaultdict(float)  # π 初始概率
    trans_mat: dict = defaultdict(lambda: defaultdict(float))  # 状态转移概率
    emission_prob: dict = defaultdict(lambda: defaultdict(float))  # 发射概率(观测概率)

    def __init__(
        self,
        input_path: str,
        output_path: str = "output.json",
        verbose: bool = False,
        print_every: int = 10000,
    ) -> None:
        """初始化分词训练器

        Args:
            input_path (str): 输入文件路径
            output_path (str, optional): 输出文件路径，json格式 . Defaults to "output.json".
            verbose (bool, optional): 是否打印信息. Defaults to False.
            print_every (int, optional): 打印的话，处理多少行后打印一次. Defaults to 10000.
        """

        self.input_path = input_path
        self.output_path = output_path
        self.verbose = verbose
        self.print_every = print_every

    def train(self):
        """开始训练，先基于字典保存"""

        # 对概率值取对数， 连乘变成连加
        A_counter = defaultdict(lambda: defaultdict(int))  # 由状态i转移到状态j的频数
        B_counter = defaultdict(lambda: defaultdict(int))  # 状态为j观测为k的频数
        init_counter = defaultdict(int)  # 状态i作为初始状态的频数

        line_num = 0  # 当前处理的行数
        seen_tokens = set()  # 已保存的token

        with open(self.input_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        # 注意，每行相当于一个样本，而不是每个单词
        for line in lines:
            # 可能的问题 数字和中文连在一起了，比如 6年 1日
            line_num += 1

            # 处理每行
            # 先去掉前后空白符
            line = line.strip()
            # 如果去掉空格后变成了空行
            if not line:
                continue

            # 根据空白符分隔
            tokens = line.split()
            # 加入到token集合
            seen_tokens.update(tokens)
            # 添加对应的状态 BMES
            states = []
            for token in tokens:
                if len(token) == 1:
                    # 单字词，对应S
                    states.append("S")
                else:
                    states.append("B" + "M" * (len(token) - 2) + "E")

            # 现在有了token，以及对应的状态： 比如 [我 是 中国人] [S S BME]
            # 但是此时需要将它们拼起来，变成 我是中国国人 SSBME
            token_line = "".join(tokens)
            state_line = "".join(states)
            # 它们俩的长度应该相等
            assert len(token_line) == len(state_line)
            for i in range(len(token_line)):
                # 初始状态
                if i == 0:
                    init_counter[state_line[i]] += 1
                else:
                    # 由i-1状态变成i状态的频次
                    # 更新A
                    A_counter[state_line[i - 1]][state_line[i]] += 1
                    # 更新B 状态为i观测为字char的频数
                B_counter[state_line[i]][token_line[i]] += 1

            if self.verbose and line_num % self.print_every == 0:
                print(line_num)

        print(f"Total tokens size: {len(seen_tokens)}.")
        print(
            "Complete the calculation of A and B, and now start calculating their log probabilities."
        )

        # 计算初始概率
        # 所有初始状态的频数之和
        total_init_chars = sum(init_counter.values())
        for state, value in init_counter.items():
            # 所有的概率计算都取对数，使得后面的概率连乘变成连加
            self.init_prob[state] = log(value / total_init_chars)

        # 计算状态转移概率
        for state, state_dict in A_counter.items():
            # state_dict 是当前状态 state 转移到其他状态(包括自己)的频次
            total_state_count = sum(state_dict.values())
            for s, v in state_dict.items():
                self.trans_mat[state][s] = log(v / total_state_count)

        # 计算发射概率
        for state, state_dict in B_counter.items():
            # state_dict 是当前状态 state 观测到不同char的频次
            total_char_count = sum(state_dict.values())
            for ch, value in state_dict.items():
                self.emission_prob[state][ch] = log(value / total_char_count)

        data = {
            "init_prob": self.init_prob,
            "trans_mat": self.trans_mat,
            "emission_prob": self.emission_prob,
            "chars": list(
                {char for token in seen_tokens for char in token}
            ),  # 所有训练数据中已知的字(char) ，set不能用于JSON序列号，先转换为set去重后，再转换为list
        }

        with open(self.output_path, "w", encoding="utf-8") as f:
            # 计算完成，保存
            json.dump(data, f)

        print(f"Success save data to {self.output_path}")


if __name__ == "__main__":
    trainer = TokenizerTrainer("RenMinData.txt_utf8", verbose=True)
    trainer.train()
